import os
import numpy as np
import pickle

class DictionaryMap(object):

    @staticmethod
    def dictionary_from_txt_files(files):
        print('正在生成字典......, 返回参数：(word2ids, ids2word)')
        w2i_dict = {'bos': 0, 'eos': 1}
        i2w_dict = {}
        for file in files:
            file_name = os.path.join("data", file)
            with open(file_name, encoding='utf8') as f:
                fline = f.readline()
                while fline:
                    for w in fline.replace('\n', '').split():
                        if w not in w2i_dict:
                            w2i_dict[w] = np.int32(len(w2i_dict))
                    fline = f.readline()
        i2w_dict = {i: w for w, i in w2i_dict.items()}
        #print('w2i_dict:', w2i_dict)
        return w2i_dict, i2w_dict

    @staticmethod
    def dictionary_from_pkl_file(filename, if_file_in_data=True):
        if if_file_in_data:
            filename = os.path.join("data", filename)
        print('从%s读取字典 ......,返回参数：(word2ids, ids2word)' % filename)
        with open(filename, "rb") as f:
            w2i, i2w = pickle.load(f)
        return w2i, i2w

    @staticmethod
    def dictionary_2_pkl_file(dictionary, filename, if_file_in_data=True):
        #if if_file_in_data:
            #filename = os.path.join("data", filename)
        print('Write dictionary to %s'%filename)
        with open(filename, 'wb',encoding='utf8') as w:
            pickle.dump(dictionary, w)